# -*- coding: utf-8 -*-
"""
@Time: 2020/12/14 20:00
@Author: 鹄望潇湘
@File: FCN_VGG16_Tester.py
@desc: FCN_VGG16网络效果测试器，用来展示FCN_VGG16网络的预测效果。
"""

from FCN_VGG16 import FCN_VGG16
import torch
import pickle
import numpy as np
from PIL.Image import Image
from torchvision.transforms import Compose
from torchvision import transforms
from PascalDataSet import ImageScaleCrop, PascalDataSetGenerator
from tools_zsj.ZsjDataSet import SingleImageDataSet, MultiDataSet
from torch.utils.data import DataLoader
from torch import Tensor
import threading


class PixelEvaluator(object):
    """

    :param class_number: 像素所属类别的总数
    """

    def __init__(self, class_number: int):
        self.class_number = class_number
        self.data_metric = np.zeros(shape=(class_number, class_number), dtype=np.int64)
        self.matrix_lock = threading.Lock()

    def push_one(self, predict_data: Tensor, real_data: Tensor):
        """

        :param predict_data: 网络的预测值，应该是一个维的矩阵
        :param real_data: 对应的真实值，应该是一个三维矩阵
        :return:
        """
        temp_matrix = np.zeros(shape=(self.class_number, self.class_number), dtype=np.int64)
        length = real_data.shape[0] * real_data.shape[1]
        real_data = real_data.reshape(length).to('cpu').numpy()
        predict_data = predict_data.reshape(length).to('cpu').numpy()

        for index in range(0, length):
            real_value = real_data[index]
            predict_value = predict_data[index]
            temp_matrix[real_value, predict_value] += 1
        pass

        self.matrix_lock.acquire()
        self.data_metric = np.add(self.data_metric, temp_matrix)
        self.matrix_lock.release()

    def push_batch(self, predict_data: Tensor, real_data: Tensor):
        """

        :param predict_data: 网络的一组预测值， 应该是一个四维矩阵，第一维表示batch_size
        :param real_data: 对应的一组真实值， 应该是一个三矩阵，第一维表示batch_size
        :return:
        """
        threads = []
        for index in range(0, real_data.shape[0]):
            try:
                th = threading.Thread(target=self.push_one, args=[predict_data[index], real_data[index]])
                threads.append(th)
                th.start()
            except Exception:
                print("线程启动失败")
        for item in threads:
            item.join()


    def show(self):
        print("总类别数:%d" % self.class_number)
        print("预测正确的像素数/总像素数: %d/%d" % (self.get_correct_sum(), self.get_pixel_sum()))
        print("mean accuracy: %.3f" % self.get_mean_accuracy())
        print("mean IU: %.3f" % self.get_mean_IU())
        print("frequency weighted IU: %.3f" % self.get_frequency_weighted_IU())

    def get_ti(self, i: int):
        return np.sum(self.data_metric[i])

    def get_correct_sum(self) -> int:
        """
        计算预测正确的像素值总数

        :return:
        """
        results = 0
        for index in range(0, self.class_number):
            results += self.data_metric[index][index]
        return results

    def get_correct_accuracy(self):
        return (self.get_correct_sum()+0.0)/self.get_pixel_sum()

    def get_mean_accuracy(self):
        """
        计算 mean_accuracy

        :return:
        """
        results = 0.0
        for index in range(0, self.class_number):
            temp = (self.get_ti(index)+0.0)
            if temp != 0:
                results += self.data_metric[index][index]/temp

        return results/self.class_number

    def get_mean_IU(self) -> float:
        """
        计算mean IU

        :return:
        """
        results = 0.0
        for index in range(0, self.class_number):
            denominator = (self.get_ti(index)+np.sum(self.data_metric[:, index]) -
                           self.data_metric[index][index])+0.0
            if denominator != 0:
                results += self.data_metric[index][index]/denominator

        return results/self.class_number

    def get_frequency_weighted_IU(self) -> float:
        """
        计算 frequency weighted IU

        :return:
        """
        results = 0.0
        sum_tk = 0.0
        for index in range(0, self.class_number):
            ti = self.get_ti(index)
            sum_tk += ti

            denominator = (ti + np.sum(self.data_metric[:, index]) -
                           self.data_metric[index][index]) + 0.0
            if denominator != 0:
                results += ti*self.data_metric[index][index]/denominator
            else:
                results += 1
        return results/sum_tk

    def get_pixel_sum(self):
        return np.sum(self.data_metric)

    def show_matrix(self):
        print(self.data_metric)


class FCN_VGG16_Tester(object):
    """
    FCN_VGG16网络效果测试器，可以用来对输入图片进行预测并产生可视化图片输出，也提供接口在特定数据集上进行模型评估（如计算交并比等）

    :param model_dict_path: 保存的模型参数文件的路径地址
    :param color_object_dict_path: 颜色与类别对应关系的字典文件路径(color_object_dict.pickle)
    :param devices: 网络计算使用的设备(cpu，cuda0等)
    """

    def __init__(self, model_dict_path: str, color_object_dict_path: str, devices):
        self.net = FCN_VGG16(init_weights=False)
        self.devices = devices
        self.net.to(devices)

        self.model_dict_path = model_dict_path
        self.color_object_dict_path = color_object_dict_path

        self.__restore_model()

        self.color_object_dict = dict()
        self.object_color_dict = dict()

        self.read_color_object_dict()
        self.create_object_color_dict()

        self.input_transformer = Compose([ImageScaleCrop(32),
                                          transforms.ToTensor()])
        self.output_transformer = Compose([transforms.ToTensor(),
                                           transforms.ToPILImage()])

    def __restore_model(self):
        self.net.load_state_dict(torch.load(self.model_dict_path, map_location=self.devices))

    def read_color_object_dict(self):
        with open(self.color_object_dict_path, 'rb') as f:
            self.color_object_dict = pickle.load(f)
        pass

    def predict_one_image(self, image_url: str) -> Image:
        """
        输入一张图片，返回神经网络的预测输出
        :param image_url: 输入图片的路径
        :return: 网络的预测输出
        """
        dataset = SingleImageDataSet()
        dataset.set_image_transformer(self.input_transformer)
        dataset.push_single_image(image_url)
        loader = DataLoader(dataset, batch_size=1, num_workers=0, shuffle=False)

        for index, item in enumerate(loader):
            image = transforms.ToPILImage()(item[0])
            image.show()
            image.save("./images/peach_fill.jpg")
            predict = self.net(item.to(self.devices))
            """"""
            # k = predict.to('cpu').detach().numpy()
            predict = torch.transpose(predict, 0, 2)
            predict = torch.transpose(predict, 1, 3)

            predict = (torch.max(predict, dim=3)[1]).to('cpu').numpy()

            # predict = predict.reshape(predict.shape[0], predict.shape[1], 1)

            predict = np.apply_along_axis(self.__object_to_color, 2, predict)

            return self.output_transformer(predict)
        pass

    def analyse_one(self, real_image_url: str, label_image_url) -> PixelEvaluator:
        """
        对一张图片的预测结果进行评估

        :param real_image_url:原图片存放路劲
        :param label_image_url:标注图片存放路径
        :return:
        """
        generator = PascalDataSetGenerator(self.color_object_dict_path)
        dataset = generator.create_dataset_by_image(real_image_url, label_image_url)
        return self.analyse_datasets(dataset)

    def analyse_datasets(self, dataset: MultiDataSet) -> PixelEvaluator:
        """
        对一个数据集进行预测并进行评估
        :param dataset: 数据集，里面应该包含原图片和预测图片
        :return:
        """
        results = PixelEvaluator(len(self.color_object_dict))
        loader = DataLoader(dataset, batch_size=1, num_workers=0, shuffle=False)
        for index, item in enumerate(loader):
            print("预测第%d张图片..." % index)
            image_data = item[0]
            label_data = item[1]

            predict = self.net(image_data.to(self.devices))
            k = predict.to('cpu').detach().numpy()
            predict = torch.max(predict, dim=1)[1]
            results.push_batch(predict, label_data.to(self.devices))
            if (index+1) % 50 == 0:
                break

        return results

    def __object_to_color(self, vec):
        return self.object_color_dict[vec[0]]

    def create_object_color_dict(self):
        for key in self.color_object_dict.keys():
            array = np.array([key[0], key[1], key[2]], dtype=np.uint8)
            m = str(array)
            if str(array) != "[128 192   0]":
                array = np.array([0, 0, 0], dtype=np.uint8)
            else:
                array = np.array([255, 255, 255], dtype=np.uint8)
            self.object_color_dict[self.color_object_dict[key]] = array
        pass


if __name__ == '__main__':
    device = torch.device(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
    tester = FCN_VGG16_Tester("./checkpoint7/checkpoint_best_7.pth", "./color_object_dict.pickle", device)

    original_image_path = "/home/xiaoxiang/DataSet/VOC2012/JPEGImages"
    labeled_image_path = "/home/xiaoxiang/DataSet/VOC2012/SegmentationClass"

    generator = PascalDataSetGenerator(color_object_dict_path="color_object_dict.pickle")
    """
    val_dataset = generator.create_dataset_by_txt(
        "/home/xiaoxiang/DataSet/VOC2012/ImageSets/Segmentation/val.txt",
        original_image_path,
        labeled_image_path)
    
    print("开始预测...")
    evaluator = tester.analyse_datasets(val_dataset)
    print("开始评估....")
    evaluator.show()
    evaluator.show_matrix()

    """
    # result = tester.predict_one_image("/home/xiaoxiang/桌面/test.jpg")
    result = tester.predict_one_image("./images/peach.jpg")
    result.show()
    result.save("./images/divided.png")

    """
    evaluator = tester.analyse_one("./images/2007_007836.jpg",
                                   "./images/2007_007836.png")
    
    evaluator.show()
    evaluator.show_matrix()
    """

